#include "train.hpp"
#include "common.h"
#include "environment.hpp"
#include "route.hpp"
#include <cassert>
#include <cstdio>
#include <iostream>
#include <numeric>
#include <vector>

Train::Train(std::vector<Locomotive *> locomotives,
             std::vector<RailroadWagon *> wagons) {
    for (auto loco : locomotives) {
        this->locomotives.push_back(loco);
    }
    for (auto wagon : wagons) {
        this->wagons.push_back(wagon);
    }
    this->route = nullptr;
    // 计算火车的总重量
    fp_type trainMass = 0;
    trainMass += std::accumulate(
        this->locomotives.begin(), this->locomotives.end(), 0.0,
        [](fp_type sum, Locomotive *loco) { return sum + loco->mass(); });
    trainMass += std::accumulate(this->wagons.begin(), this->wagons.end(), 0.0,
                                 [](fp_type sum, RailroadWagon *wagon) {
                                     return sum + wagon->totalMass();
                                 });

    this->trainDynamics.mass = trainMass;
    // 初始化其它字段
    this->trainDynamics.state = STOP;
    this->trainDynamics.length = calcLength();
    this->trainDynamics.speed = 0.0;
    this->trainDynamics.distance = 0.0;
    this->trainDynamics.acceleration = 0.0;
    this->trainDynamics.direction = DIRECTION::DOWN;
    this->trainDynamics.currCycPwrCons = 0.0;
}

void Train::setRoute(Route *route) {
    // 配置火车的路线
    this->route = route;
    // 根据火车的长度，确定车头的位置，假设车尾在起始位置。
    assert(this->trainDynamics.length <
           route->getCurrentRouteItem()->track->getLength());
    this->route->headPosition = this->trainDynamics.length;
    this->route->tailPosition = 0;
}

fp_type Train::calcLength() {
    fp_type length = 0.0;
    length += std::accumulate(
        this->locomotives.begin(), this->locomotives.end(), (fp_type)0.0,
        [](fp_type sum, Locomotive *loco) { return sum + loco->getLength(); });
    length +=
        std::accumulate(this->wagons.begin(), this->wagons.end(), (fp_type)0.0,
                        [](fp_type sum, RailroadWagon *wagon) {
                            return sum + wagon->getLength();
                        });
    return length;
}

fp_type Train::maxTraction() {

    fp_type totalTraction = 0.0;
    fp_type speed = this->trainDynamics.speed;
    totalTraction +=
        std::accumulate(this->locomotives.begin(), this->locomotives.end(),
                        (fp_type)0.0, [speed](fp_type sum, Locomotive *loco) {
                            return loco->traction(loco->maxPower, speed, 0.0);
                        });
    return totalTraction;
}

fp_type Train::totalFriction() {
    fp_type totalFriction = 0.0;
    totalFriction +=
        std::accumulate(this->locomotives.begin(), this->locomotives.end(),
                        (fp_type)0.0, [](fp_type sum, Locomotive *loco) {
                            return sum + loco->friction(0.0, 0.0);
                        });
    totalFriction +=
        std::accumulate(this->wagons.begin(), this->wagons.end(), (fp_type)0.0,
                        [](fp_type sum, RailroadWagon *wagon) {
                            return sum + wagon->friction(0.0, 0.0);
                        });
    // for (auto loco : this->locomotives) {
    //     totalFriction += loco->friction(this->trainDynamics.speed, 0);
    // }
    // for (auto wagon : this->wagons) {
    //     totalFriction += wagon->friction(this->trainDynamics.speed, 0);
    // }
    return totalFriction;
}

fp_type Train::calcAcceleration() {
    fp_type trainMaxTraction = this->maxTraction();
    fp_type totalFriction = this->totalFriction();
    return (trainMaxTraction - totalFriction) * 1000 /
           (this->trainDynamics.mass * 1000);
}

// 根据当前的阻力和加速度，反推出当前功率
// 单位：kW
fp_type Train::getPower() {
    const fp_type accel = this->trainDynamics.acceleration;
    // 总牵引力，单位kN
    const fp_type tractionForce =
        this->totalFriction() +
        (this->trainDynamics.mass * 1000 * accel / 1000);
    // 电能消耗，单位kW
    return tractionForce * 1000 * this->trainDynamics.speed / 3.6 / 1000;
}
// 更新运行状态
void Train::runningStatusStep() {
    if (this->route == nullptr)
        return;

    RouteItem *currentRouteItem = this->route->getCurrentRouteItem();
    if (currentRouteItem == nullptr) {
        std::cout << "route is end!!" << std::endl;
        return;
    }
    // 每5s执行一次区间是否变更的判断
    if (environment.t % (5) == 0) {

        // 如果车头在当前区间中驶过的里程大于该区间的长度，那么就可以将车头的区间序号增加一个。
        if (currentRouteItem->track->getLength() <= this->route->headPosition) {
            this->route->headPosition = 0;
            this->route->incHeadIndex();
        };
        RouteItem *currentRouteTailItem =
            this->route->getCurrentTailRouteItem();
        if (currentRouteTailItem == nullptr) {
            std::cout << "route is end!!" << std::endl;
            // TODO: 如果离开最后一个区间时的速度过快，则可认为是发生了事故。
            this->trainDynamics.state = TRAIN_STATE::STOP;
            return;
        }
        // 如果车尾在当前区间中驶过的里程大于该区间的长度，那么就可以将车尾的区间序号增加一个。

        if (currentRouteTailItem->track->getLength() <=
            this->route->tailPosition) {
            this->route->tailPosition = 0;
            this->route->incTailIndex();
        }

        // 刹车距离估算，单位为m，设置其减速度为1.0m/s^2
        // x = v0^2 / (2*a)
        // 同时增加5s当前速度走过的距离作为缓冲距离
        fp_type brakeDistance =
            (this->trainDynamics.speed * this->trainDynamics.speed) /
                (2.0 * 1.0 * 3.6 * 3.6) +
            5 * this->trainDynamics.speed / 3.6;
        std::cout << "distance to terminal: "
                  << this->route->getDistanceToTerminal() << std::endl;
        std::cout << "brake distance: " << brakeDistance << std::endl;
        // 判断与终点的距离
        // 还要判断前后列车状态
        if (this->route->getDistanceToTerminal() <= brakeDistance) {
            this->trainDynamics.state = TRAIN_STATE::BRAKING;
            this->trainDynamics.targetSpeed = 20; // 刹车到低速接近站台
        } else {
            this->trainDynamics.targetSpeed =
                this->route->getCurrentRouteItem()->track->getSpeedLimit();
        }
    }
    // 计算状态转移，并更新加速度和速度
    switch (this->trainDynamics.state) {
    case WAITING:
        this->trainDynamics.acceleration = 0;
        this->trainDynamics.speed = 0;
        break;
    case ACCELERATING:
        this->trainDynamics.acceleration = this->calcAcceleration();
        this->trainDynamics.speed +=
            (this->trainDynamics.acceleration * 3.6) * environment.dt;
        if (this->trainDynamics.targetSpeed - 5.0 <=
            this->trainDynamics.speed) {
            this->trainDynamics.state = TRAIN_STATE::EVEN_SPEED;
        }
        break;
    case EVEN_SPEED:
        this->trainDynamics.acceleration = 0.0;
        if (this->trainDynamics.speed >=
            (this->trainDynamics.targetSpeed + 3.0)) {
            this->trainDynamics.state = TRAIN_STATE::BRAKING;
        } else if (this->trainDynamics.speed <=
                   (this->trainDynamics.targetSpeed - 3.0)) {
            this->trainDynamics.state = TRAIN_STATE::ACCELERATING;
        }
        break;
    case BRAKING:
        this->trainDynamics.acceleration = -1.0 * environment.dt;
        this->trainDynamics.speed +=
            this->trainDynamics.acceleration * 3.6 * environment.dt;
        if (this->trainDynamics.speed >= 1.0 &&
            this->trainDynamics.speed <= this->trainDynamics.targetSpeed) {
            this->trainDynamics.state = TRAIN_STATE::EVEN_SPEED;
            this->trainDynamics.acceleration = 0.0;
        } else if (this->trainDynamics.speed < 1.0) {
            this->trainDynamics.state = TRAIN_STATE::STOP;
            this->trainDynamics.speed = 0.0;
            this->trainDynamics.acceleration = 0.0;
        }
        break;
    default:
        break;
    }
    fp_type delta_distance = 0.0;
    delta_distance = this->trainDynamics.speed / 3.6 * environment.dt;
    this->route->headPosition += delta_distance;
    this->route->tailPosition += delta_distance;
    this->trainDynamics.distance += delta_distance;
}
/**
每一个仿真步中，火车状态的更新。
火车状态管理
*/

void Train::update() {

    // 需要梳理这里的状态，外面的大状态用if-else判断，
    // 里面的小状态用switch-case判断
    if (IS_TRAIN_STATE_RUNNING(this->trainDynamics.state)) {
        this->runningStatusStep();
    }

    // 计算当前周期内的能量消耗
    const fp_type energy = this->getPower() * environment.dt;
    printf("train targetSpeed %lf speed %lf, power %lf, distance %lf, head "
           "%d:%lf, "
           "tail %d:%lf \n",
           (fp64)this->trainDynamics.targetSpeed,
           (fp64)this->trainDynamics.speed, (fp64)this->getPower(),
           (fp64)this->trainDynamics.distance, this->route->getHeadIndex(),
           (fp64)this->route->headPosition, this->route->getTailIndex(),
           (fp64)this->route->tailPosition);
}
